# pip install langchain pymilvus

from langchain.vectorstores import Milvus
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document

# ========================
# 1️⃣ 云端 Milvus 配置
# ========================
MILVUS_COLLECTION = "test_collection"
ZILLIZ_ENDPOINT = "https://in03-5451ecee2d84db2.serverless.aws-eu-central-1.cloud.zilliz.com"
ZILLIZ_API_KEY = "378d8b31d9aacf110c136431202de4613e31832d43786bffd255d4f13a725f2fdd1f81fe84c62fd3adb038f6bf3b70673b582bd5"  # 从 Zilliz 控制台获取的 Token/API Key

#
# ========================
# 2️⃣ 初始化嵌入模型（本地或云端都可以）
# ========================
embedding_model = HuggingFaceEmbeddings(
    model_name="D:\\models\\models\\sentence-transformers\\paraphrase-multilingual-MiniLM-L12-v2"
)

# ========================
# 3️⃣ 初始化 Milvus 向量库对象
# ========================
db = Milvus(
    embedding_function=embedding_model,
    collection_name=MILVUS_COLLECTION,
    # 针对 Zilliz Cloud 的配置
    connection_args={
        "uri": ZILLIZ_ENDPOINT,
        "token": ZILLIZ_API_KEY,  # Zilliz Cloud 的 API Key 或 Token
        "secure": True,  # 启用 SSL/TLS
    },
    auto_id=True  # ✅ 自动生成 id
)

# ========================
# 4️⃣ 增加文档（可批量）
# ========================
docs = [
    Document(page_content="管理学是一门研究企业和组织管理的学科。", metadata={"id": "doc1"}),
    Document(page_content="贝尔宾的团队角色理论可以帮助分析团队成员特点。", metadata={"id": "doc2"}),
]
db.add_documents(docs)
print("文档已加入 Milvus")

# ========================
# 5️⃣ 相似度检索
# ========================
query = "贝尔宾团队角色"
result = db.similarity_search(query, k=3)  # 返回最相似的 3 条
for doc in result:
    print("检索结果:", doc.page_content)
